import os
import pandas as pd
import numpy as np
from sklearn.metrics import mean_squared_error, mean_absolute_error, log_loss, roc_auc_score, recall_score, r2_score
from get_data_and_prompt import (
  get_train_test_data
)
from consts import (
  FEATURE_KEYWORD,
  FEATURE_GENRE,
  FEATURE_INFO,
  FEATURE_ALL,
)
from utils import (
  log_message
)




def get_score(feature_type):
  input_directory_path = f'../data/result/basic_large_model/final_data/{feature_type}'
  output_file_path = f'../data/result/basic_large_model/result_for_{feature_type}.csv'
  data = get_train_test_data(feature_type)
  test_data = data['test_data']
  # 结果集合
  result = []
  # 遍历目录下的所有文件
  for filename in os.listdir(input_directory_path):
    # 检查文件是否为CSV文件
    if filename.endswith('.csv'):
      log_message(f"文件名称{filename}")
      file_path = os.path.join(input_directory_path, filename)
      result_data = pd.read_csv(file_path)

      merged_df = pd.merge(test_data, result_data, on='userId', how='outer')
      # 找出不同的结果
      different_data = merged_df[merged_df['rating'] != merged_df['predict']]
      # 均方误差
      mse = mean_squared_error(test_data['rating'], result_data['predict'])
      rmse = np.sqrt(mse)  # 计算RMSE，即MSE的平方根
      # 平均绝对误差
      mae = mean_absolute_error(test_data['rating'], result_data['predict'])
      log_message(f"RMSE:{rmse}")
      log_message(f"MAE:{mae}")
      # 模型名称
      model_name = filename.split('_')[0]
      result.append({
        "model": model_name,
        "rmse": round(rmse, 4),
        "mae": round(mae, 4),
      })
  result_df = pd.DataFrame(result)
  result_df.to_csv(output_file_path, index=False)
  log_message("结果存储完毕")

def main():
  get_score(FEATURE_KEYWORD)
  get_score(FEATURE_GENRE)
  get_score(FEATURE_INFO)
  get_score(FEATURE_ALL)

main()

